package com.codeing.demo.utils;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.GeneralSecurityException;
import java.security.SecureRandom;

/**
 * AES加密
 */
public class AESUtil {
    private static Logger LOGGER = LoggerFactory.getLogger(AESUtil.class);

    private static final int DEFAULT_AES_KEYSIZE = 128;

    private static final int DEFAULT_IVSIZE = 16;

    private static SecureRandom random = new SecureRandom();

    private String transformation;

    private SecretKey key;

    private IvParameterSpec iv;

    private AESUtil(byte[] keyBytes, byte[] ivBytes, String transformation, int keyBitsSize){
        this.transformation = transformation;
        byte[] keyBytes2 = new byte[keyBitsSize / 8];
        for (int i = 0; i < keyBytes2.length && i < keyBytes.length; i++) {
            keyBytes2[i] = keyBytes[i];
        }
        this.key = new SecretKeySpec(keyBytes2, "AES");
        if (!transformation.contains("ECB")) {
            byte[] ivBytes2 = new byte[16];
            for (int i = 0; i < ivBytes2.length && i < ivBytes.length; i++) {
                ivBytes2[i] = ivBytes[i];
            }
            this.iv = new IvParameterSpec(ivBytes2);
        }
    }

    /**
     * AES/CBC/PKCS5Padding(128)
     * @param keyBytes
     * @param ivBytes
     * @return
     */
    public static AESUtil getInstance$CBC_PKCS5Padding(byte[] keyBytes, byte[] ivBytes) {
        try {
            return new AESUtil(keyBytes, ivBytes, "AES/CBC/PKCS5Padding", 128);
        } catch (Exception e) {
            LOGGER.error("", e);
            return null;
        }
    }

    /**
     * AES/ECB/NoPadding(128)
     * @param keyBytes
     * @return
     */
    public static AESUtil getInstance$ECB_PKCS5Padding(byte[] keyBytes) {
        try {
            return new AESUtil(keyBytes, null, "AES/ECB/PKCS5Padding", 128);
        } catch (Exception e) {
            LOGGER.error("", e);
            return null;
        }
    }


    public byte[] decrypt(byte[] src) {
        byte[] ret = null;
        try {
            Cipher cipher = Cipher.getInstance(transformation);
            if (!this.transformation.contains("ECB")) {
                cipher.init(Cipher.DECRYPT_MODE, this.key, this.iv);
            } else {
                cipher.init(Cipher.DECRYPT_MODE, this.key);
            }
            ret = cipher.doFinal(src);
        } catch (Exception ex) {
            ex.printStackTrace();
            LOGGER.error("解密失败", ex);
        }
        return ret;
    }

    public byte[] encrypt(String src) {
        byte[] ret = null;
        try {
            Cipher cipher = Cipher.getInstance(transformation);
            if (!this.transformation.contains("ECB")) {
                cipher.init(Cipher.ENCRYPT_MODE, this.key, this.iv);
            } else {
                cipher.init(Cipher.ENCRYPT_MODE, this.key);
            }
            ret = cipher.doFinal(src.getBytes());
        } catch (Exception ex) {
            LOGGER.error("加密失败", ex);
        }
        return ret;
    }

    /**
     * 生成AES密钥,返回字节数组, 默认长度为128位(16字节).
     */
    public static byte[] generateAesKey() {
        return generateAesKey(DEFAULT_AES_KEYSIZE);
    }

    /**
     * 生成AES密钥,可选长度为128,192,256位.
     */
    public static byte[] generateAesKey(int keysize) {
        try {
            KeyGenerator keyGenerator = KeyGenerator.getInstance("AES");
            keyGenerator.init(keysize);
            SecretKey secretKey = keyGenerator.generateKey();
            return secretKey.getEncoded();
        } catch (GeneralSecurityException e) {
            LOGGER.error("生成密钥失败", e);
            return null;
        }
    }

    /**
     * 生成随机向量,默认大小为cipher.getBlockSize(), 16字节.
     */
    public static byte[] generateIV() {
        byte[] bytes = new byte[DEFAULT_IVSIZE];
        random.nextBytes(bytes);
        return bytes;
    }
}
